import torch
from torch import nn
import math
from torch.nn import functional as F
class SelfAttention(nn.Module):
    def __init__(self, head_num, dim_embed, in_proj_bias=True, out_proj_bias=True):
        super(SelfAttention,self).__init__()
        self.head_num = head_num
        self.dim_embed = dim_embed
        self.dim_per_head = self.dim_embed // self.head_num
        self.in_proj = nn.Linear(dim_embed, 3 * dim_embed, bias=in_proj_bias)
        self.out_proj = nn.Linear(dim_embed, dim_embed, bias=out_proj_bias)

    def forward(self, x, causal_mask=False):
        input_shape = batch_size, seq_length, dim_embed = x.shape
        interim_shape = (batch_size, seq_length, self.head_num, self.dim_per_head)
        q, k, v = self.in_proj(x).chunk(3,dim=-1)
        q = q.view(interim_shape).transpose(1, 2)
        k = k.view(interim_shape).transpose(1, 2)
        v = v.view(interim_shape).transpose(1, 2)
        score = q @ k.transpose(2, 3) / math.sqrt(self.dim_per_head)
        if causal_mask:
            mask = torch.ones_like(score, dtype=bool).triu(1) 
            score = score.masked_fill_(mask,-torch.inf)
        score = F.softmax(score, dim=-1) @ v
        score = score.transpose(1, 2)
        output = score.reshape(input_shape)
        return self.out_proj(output)

class CrossAttention(nn.Module):
    def __init__(self, head_num, dim_embed, dim_cross, in_proj_bias=True, out_proj_bias=True):
        super().__init__()
        self.q_proj = nn.Linear(dim_embed, dim_embed, bias=in_proj_bias)
        self.k_proj = nn.Linear(dim_cross, dim_embed, bias=in_proj_bias)
        self.v_proj = nn.Linear(dim_cross, dim_embed, bias=in_proj_bias)
        self.out_proj = nn.Linear(dim_embed, dim_embed, bias=out_proj_bias)
        self.head_num = head_num
        self.dim_head= dim_embed // self.head_num

    def forward(self, x, y):
        #x:latent:(batch_size, seq_len_q, dim_q)
        #y:context:(batch_size, seq_len_kv, dim_kv)
        input_shape = x.shape
        batch_size, seq_length, dim_embed = input_shape
        interim_shape = (batch_size, -1, self.head_num, self.dim_head)

        q = self.q_proj(x)
        k = self.k_proj(y)
        v = self.v_proj(y)
        q = q.view(interim_shape).transpose(1, 2)
        k = k.view(interim_shape).transpose(1, 2)
        v = v.view(interim_shape).transpose(1, 2)

        score = q @ k.transpose(-1, -2)
        score /= math.sqrt(self.dim_head)
        score = F.softmax(score, dim=-1)
        output = score @ v
        output = output.transpose(1, 2).contiguous().view(input_shape)
        return self.out_proj(output)
    
